#!/usr/bin/env python
# -*- conding: utf-8 -*-

"""
@Time     : 2024/9/12 6:12
@Author   : liujingmao
@File     : 1.少量示例提示模板.py
"""

import dotenv
from langchain_core.output_parsers import StrOutputParser

from langchain_core.prompts import ChatPromptTemplate, FewShotChatMessagePromptTemplate
from langchain_openai import ChatOpenAI

dotenv.load_dotenv()

"""
# 1. 构建示例模板与示例
# 2. 构建少量示例提示模板
# 3. 构建最终提示模板
# 4. 创建大语言模型与链
# 5. 调用链结果

"""

# 1. 构建示例模板与示例
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{question}"),
        ("ai", "{answer}"),
    ]
)

examples = [
    {"question": "帮我计算下2+2等于多少？", "answer": "7"},
    {"question": "帮我计算下2+3等于多少？", "answer": "5"},
    {"question": "帮我计算下20*15等于多少？", "answer": "300"},
]

# 2. 构建少量示例提示模板
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples
)
print("少量示例模板", few_shot_prompt.format())

# 3. 构建最终提示模板
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", "你是一个可以计算复杂数学问题的聊天机器人"),
        few_shot_prompt,
        ("human", "{question}"),
    ]
)
# 4. 创建大语言模型与链
llm = ChatOpenAI(model="gpt-3.5-turbo-16k", temperature=0)
chain = prompt | llm | StrOutputParser()
# 5. 调用链结果
print(chain.invoke("帮我计算一下14乘以15等于多少"))

print(chain.invoke("帮我计算一下2加2等于多少"))

"""
少量示例模板 Human: 帮我计算下2+2等于多少？
AI: 7
Human: 帮我计算下2+3等于多少？
AI: 5
Human: 帮我计算下20*15等于多少？
AI: 300
14乘以15等于210。
2加2等于4。

"""
